import struct

import torch
from torch.utils.data import Dataset
import numpy as np
import os
class Voxels(object):
    def __init__(self, data, dims, translate, scale, axis_order):
        self.data = data
        self.dims = dims
        self.translate = translate
        self.scale = scale
        # axis不是xyz坐标系的时候触发异常
        assert (axis_order in ('xzy', 'xyz'))
        self.axis_order = axis_order

def read_header(fp):
    # 读取binvox头文件
    line = fp.readline().strip()
    if not line.startswith(b'#binvox'):
        raise IOError('Not a binvox file')
    #这里实际上对输入进来的相关信息做了正则分割
    dims = list(map(int, fp.readline().strip().split(b' ')[1:]))
    translate = list(map(float, fp.readline().strip().split(b' ')[1:]))
    scale = list(map(float, fp.readline().strip().split(b' ')[1:]))[0]
    line = fp.readline()
    return dims, translate, scale
def getNeighbour(i,dims):
    #我破防了，我没看懂这段代码是做什么用的
    #这是之前写的，现在改完之后懂了，magic number死个妈
    x = i % dims[0]
    y = int(i / dims[0]) % dims[1]
    z = int(i / dims[0] / dims[1]) % dims[2]
    p = [x, y, z]
    neighbours = []
    if p[0] > 0:
        neighbours.append([p[0]-1, p[1], p[2]])
    if p[1] > 0:
        neighbours.append([p[0], p[1]-1, p[2]])
    if p[2] > 0:
        neighbours.append([p[0], p[1], p[2]-1])
    if p[0] < dims[0]-1:
        neighbours.append([p[0]+1 , p[1] , p[2]])
    if p[1] < dims[1]-1:
        neighbours.append([p[0], p[1]+1, p[2]])
    if p[2] < dims[2]-1:
        neighbours.append([p[0], p[1], p[2]+1])
        #假设if全为真，那么返回值就最多6个向量，分别对应欧几里得距离为一的点
    return neighbours
#边缘判定
def is_edge(i,neighbours,data):
    for neighbour in neighbours:
        idx = neighbour[0]*1+neighbour[1]*128+neighbour[2]*128*128
        if((data[i]==1 and data[idx]==0) or (data[i]==0 and data[idx]==1)):
            return True
    return False
#体素文件实际上是一种特殊的形式，它将三维模型拉伸成一位数组，存储的是<占用与否，重复几次>
#其中重复几次最大为255，因为受到8进制存储的限制
def getVoxel(edges):
    ret=[]
    for edge in edges:
        if (len(ret) > 0 and edge == ret[-2] and ret[-1] < 255):
            ret[-1] += 1
        else:
            ret.append(edge)
            ret.append(1)
    return ret
def read_header2(fp):
    # 读取binvox头文件
    headline = fp.readline()
    dims = fp.readline()
    translate =fp.readline()
    scale =fp.readline()
    dataline = fp.readline()
    return headline,dims, translate, scale,dataline
class VoxDataset(Dataset):
    def __init__(self, vox_file):
        dims, translate, scale = read_header(vox_file)
        raw_data = np.frombuffer(vox_file.read(), dtype=np.uint8)
        values, counts = raw_data[::2], raw_data[1::2]
        #转成int型，为了训练
        self.data = np.repeat(values, counts).astype(np.int)

        self.dims = dims
        self.coordinateas = torch.Tensor(
            [[k, j, i]
             for i in range(self.dims[0])
             for j in range(self.dims[1])
             for k in range(self.dims[2])])
        self.translate = translate
        self.scale = scale

    def __getitem__(self, idx):
        return self.coordinateas[idx], torch.tensor(self.data[idx]).long()

    def __len__(self):
        return self.dims[0] * self.dims[1] * self.dims[2]
    #获取边缘？
    def __getedges__(self):
        edges=[]
        #这个循环会遍历所有体素块，在空间里面遍历了128*128*128次
        for i in range(self.dims[0]*self.dims[1]*self.dims[2]):
            #edges实际上是一条0 1序列，标识原有数据集中哪些是边缘，哪些不是
            neighbours = getNeighbour(i, self.dims)
            if is_edge(i, neighbours, self.data):
                edges.append(1)
            else:
                edges.append(0)
        return edges

    def __getdims__(self):
        return self.dims
class VoxedgeDataset(Dataset):
    def __init__(self, voxdata, edges):
        self.data = []
        self.datalen = 0
        for i in range(len(edges)):
            if edges[i] == 1:
                self.data.append(voxdata.__getitem__(i))
                self.datalen = self.datalen + 1

    def __getitem__(self, idx):
        return self.data[idx]
    def __len__(self):
        return self.datalen


def initdata(mypath):
    path = mypath
    with open(path, 'rb') as f:
        voxdata = VoxDataset(f)
        'print(voxdata.dims,voxdata.scale,voxdata.translate,voxdata.coordinateas)'
        'print(voxdata.data)'

        #获取体素坐标和体素是否被占用coordinateas——坐标，value——取值0或者1
        #训练时调用的核心方法，其他都别管
        #输入是第几个体素 = x * y * z
        #返回体素坐标和是否被占用
        coordinateas,value=voxdata.__getitem__(2000)
        print(coordinateas)
        print(value)
        '''
        savePath = "模型/output/5.txt"
        with open(savePath, 'w')as fp:
            for i in range(voxdata.dims[0] * voxdata.dims[1] * voxdata.dims[2]):
                coordinateas, value = voxdata.__getitem__(i)
                fp.write(str(int(coordinateas[0].item())))
                fp.write(",")
                fp.write(str(int(coordinateas[1].item())))
                fp.write(",")
                fp.write(str(int(coordinateas[2].item())))
                fp.write(",")
                fp.write(str(int(value.item())))
                fp.write("\n")

        #print(coordinateas,value)
        '''
        #所以这里时用edges给所有体素又打上了了一遍标识，是edge为1，不是为0
        #边缘再采样——不重要
        edges=voxdata.__getedges__()
        #不重要
        voxels=getVoxel(edges)

        #输出方式，一行都别改
        with open(path, 'rb') as f:
            headline, dims, translate, scale, dataline = read_header2(f)
        print(headline,dims,translate,scale,dataline)
        save_dir = "模型/output"
        save_name = "model"
        savePath = os.path.join(save_dir, save_name + ".binvox")
        with open(savePath, 'wb')as fp:
            # a = struct.pack('B', headline,dims,translate,scale,dataline)
            fp.write(headline)
            fp.write(dims)
            fp.write(translate)
            fp.write(scale)
            fp.write(dataline)
            for x in voxels:
                'print(x)'
                a = struct.pack('B', x)
                'print(a)'
                fp.write(a)
        print("done")

def savingdata(voxdata, mypath, num):
    with open(mypath, 'rb') as f:
        headline, dims, translate, scale, dataline = read_header2(f)
    print(headline, dims, translate, scale, dataline)
    #在这里修改输出的binvox文件路径和文件名
    save_dir = "模型/output"
    save_name = num.__str__() + '(final,48)'
    savePath = os.path.join(save_dir, save_name + ".binvox")
    with open(savePath, 'wb')as fp:
        # a = struct.pack('B', headline,dims,translate,scale,dataline)
        fp.write(headline)
        fp.write(dims)
        fp.write(translate)
        fp.write(scale)
        fp.write(dataline)
        for x in voxdata:
            'print(x)'
            a = struct.pack('B', x)
            'print(a)'
            fp.write(a)
    print("export data done")